import os
import re
import tempfile
import shutil
from docx import Document
import pythoncom
import win32com.client as win32
import zipfile

def extract_docx_text(file_path):
    """从DOCX文件中提取文本"""
    doc = Document(file_path)
    full_text = []
    
    # 提取段落文本
    for paragraph in doc.paragraphs:
        if paragraph.text.strip():
            full_text.append(paragraph.text)
            
    # 提取表格文本
    for table in doc.tables:
        for row in table.rows:
            for cell in row.cells:
                if cell.text.strip():
                    full_text.append(cell.text)
                    
    return "\n".join(full_text)

def extract_docx_text_enhanced(file_path):
    """增强的DOCX文本提取功能"""
    try:
        doc = Document(file_path)
        full_text = []
        
        # 提取段落文本
        for paragraph in doc.paragraphs:
            text = paragraph.text.strip()
            if text:
                full_text.append(text)
                
        # 提取表格文本
        for table in doc.tables:
            for row in table.rows:
                row_text = []
                for cell in row.cells:
                    cell_text = cell.text.strip()
                    if cell_text:
                        row_text.append(cell_text)
                if row_text:
                    full_text.append(" | ".join(row_text))
                    
        # 尝试提取页眉页脚
        try:
            for section in doc.sections:
                # 页眉
                if section.header:
                    for paragraph in section.header.paragraphs:
                        text = paragraph.text.strip()
                        if text:
                            full_text.append(f"[页眉] {text}")
                
                # 页脚
                if section.footer:
                    for paragraph in section.footer.paragraphs:
                        text = paragraph.text.strip()
                        if text:
                            full_text.append(f"[页脚] {text}")
        except:
            pass  # 忽略页眉页脚提取错误
            
        return "\n".join(full_text)
    except Exception as e:
        raise Exception(f"DOCX文件提取失败: {str(e)}")

def extract_doc_text(file_path):
    """从DOC文件中提取文本（使用Word COM接口）"""
    pythoncom.CoInitialize()  # 初始化COM
    
    # 创建临时文件
    temp_dir = tempfile.mkdtemp()
    temp_docx = os.path.join(temp_dir, "temp.docx")
    
    try:
        # 使用Word应用程序转换DOC到DOCX
        word = win32.Dispatch("Word.Application")
        word.Visible = False
        word.DisplayAlerts = False
        
        # 修复路径问题：使用原始字符串并确保路径格式正确
        # 将路径转换为绝对路径并确保使用反斜杠
        abs_path = os.path.abspath(file_path)
        # 使用原始字符串避免转义问题
        abs_path = abs_path.replace('/', '\\')
        
        # 打开文档
        doc = word.Documents.Open(abs_path)
        doc.SaveAs2(temp_docx, FileFormat=16)  # 16表示docx格式
        doc.Close()
        word.Quit()
        
        # 从转换后的DOCX中提取文本
        text = extract_docx_text(temp_docx)
        
    except Exception as e:
        # 确保在异常时也退出Word
        try:
            if 'doc' in locals() and doc:
                doc.Close()
            if 'word' in locals() and word:
                word.Quit()
        except:
            pass
        raise Exception(f"DOC文件转换失败: {str(e)}")
    finally:
        # 清理临时文件
        try:
            if 'word' in locals():
                word.Quit()
        except:
            pass
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir, ignore_errors=True)
        pythoncom.CoUninitialize()
        
    return text

def extract_doc_text_enhanced(file_path):
    """增强的DOC文本提取功能"""
    pythoncom.CoInitialize()
    
    # 直接尝试使用 Word COM 接口打开文件
    
    temp_dir = tempfile.mkdtemp()
    temp_docx = os.path.join(temp_dir, "temp.docx")
    
    try:
        # 使用Word应用程序转换DOC到DOCX
        word = win32.Dispatch("Word.Application")
        word.Visible = False
        word.DisplayAlerts = False
        
        # 处理可能的密码保护
        try:
            # 使用原始字符串路径，确保路径格式正确
            abs_path = os.path.abspath(file_path).replace('/', '\\')
            doc = word.Documents.Open(abs_path, False, True)  # 只读方式打开
        except Exception as open_error:
            # 尝试以不同方式打开
            try:
                doc = word.Documents.Open(abs_path, False, False, False, "")
            except:
                # 记录错误但继续尝试其他方法
                self.log(f"无法使用Word打开DOC文件: {str(open_error)}")
                # 尝试使用备用方法
                text = extract_text_fallback(file_path)
                if text.strip():
                    return text
                else:
                    raise Exception(f"无法打开DOC文件: {str(open_error)}")
        
        # 检查文档是否受保护
        try:
            if doc.ProtectionType != 0:  # 0表示无保护
                # 尝试取消保护（如果可能）
                try:
                    doc.Unprotect()
                except:
                    self.log(f"文档受保护，可能无法提取全部内容: {file_path}")
        except:
            pass  # 忽略保护检查错误
        
        doc.SaveAs2(temp_docx, FileFormat=16)
        doc.Close()
        word.Quit()
        
        # 从转换后的DOCX中提取文本
        text = extract_docx_text_enhanced(temp_docx)
        
        return text
        
    except Exception as e:
        # 尝试直接提取文本（作为最后手段）
        try:
            text = extract_text_fallback(file_path)
            if text.strip():
                return text
        except:
            pass
        
        raise Exception(f"DOC文件转换失败: {str(e)}")
    finally:
        # 清理临时文件
        try:
            if 'word' in locals():
                word.Quit()
        except:
            pass
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir, ignore_errors=True)
        pythoncom.CoUninitialize()

def extract_text_fallback(file_path):
    """作为最后手段的文本提取方法"""
    # 尝试使用其他方法提取文本
    try:
        # 如果是DOCX文件，尝试直接解压读取
        if file_path.endswith('.docx'):
            text_content = []
            with zipfile.ZipFile(file_path, 'r') as docx_zip:
                # 读取主要的文档内容
                if 'word/document.xml' in docx_zip.namelist():
                    with docx_zip.open('word/document.xml') as document_file:
                        content = document_file.read().decode('utf-8')
                        # 简单的XML文本提取
                        text_content.extend(re.findall(r'<w:t[^>]*>([^<]+)</w:t>', content))
                
                return ' '.join(text_content)
    except:
        pass
    
    # 如果是DOC文件，尝试使用文本提取工具
    try:
        import textract
        return textract.process(file_path).decode('utf-8')
    except:
        pass
    
    return ""

def clean_text(text, remove_empty=True, min_length=1):  # 降低最小长度限制
    """清理文本"""
    lines = text.split('\n')
    cleaned_lines = []
    
    for line in lines:
        # 移除多余空格但保留基本格式
        line = re.sub(r'\s+', ' ', line.strip())
        
        # 检查是否满足条件
        if remove_empty and not line:
            continue
            
        # 降低最小长度要求或完全移除
        if min_length > 0 and len(line) < min_length:
            # 但仍然保留可能重要的短文本（如标题）
            if any(keyword in line.lower() for keyword in ['标题', '目录', '摘要', '结论']):
                cleaned_lines.append(line)
            continue
            
        cleaned_lines.append(line)
        
    return cleaned_lines

def save_as_jsonl(text_lines, output_path):
    """保存为JSONL格式"""
    import json
    with open(output_path, 'w', encoding='utf-8') as f:
        for line in text_lines:
            json.dump({"text": line}, f, ensure_ascii=False)
            f.write('\n')
            
def save_as_txt(text_lines, output_path):
    """保存为TXT格式"""
    with open(output_path, 'w', encoding='utf-8') as f:
        for line in text_lines:
            f.write(line + '\n')
